# author: lgx
# date: 2023/4/4 9:35
# description:
import numpy as np
import torch


class Compose:
    """Compose a data pipeline with a sequence of transforms.
    Args:
        transforms (list[dict | callable]):
            Either config dicts of transforms or transform objects.
    """
    def __init__(self, transforms):
        self.transforms = []
        if isinstance(transforms, dict):
            for t, args in transforms.items():
                self.transforms.append(eval(t)(**args) if isinstance(args, dict) else eval(t)())
        elif isinstance(transforms, list):
            for transform in transforms:
                if isinstance(transform, dict):
                    transform = eval(transform['type'])(**transform)
                if callable(transform):
                    self.transforms.append(transform)
                else:
                    raise TypeError(f'transform must be callable or a dict, '
                                    f'but got {type(transform)}')

    def __call__(self, data):
        """Call function to apply transforms sequentially.
        Args:
            data (dict): A result dict contains the data to transform.
        Returns:
            dict: Transformed data.
        """
        for t in self.transforms:
            data = t(data)
            if data is None:
                return None
        return data

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        for t in self.transforms:
            format_string += '\n'
            format_string += '    {0}'.format(t)
        format_string += '\n)'
        return format_string


class Rename:
    """Rename the key in results.
    Args:
        mapping (dict): The keys in results that need to be renamed. The key of
            the dict is the original name, while the value is the new name. If
            the original name not found in results, do nothing.
            Default: dict().
    """

    def __init__(self, mapping):
        self.mapping = mapping

    def __call__(self, results):
        for key, value in self.mapping.items():
            if key in results:
                assert isinstance(key, str) and isinstance(value, str)
                assert value not in results, ('the new name already exists in '
                                              'results')
                results[value] = results[key]
                results.pop(key)
        return results


class Transpose:
    """
        Format the skeleton's shape.
        For example, transpose `SMTVC` in to `SMCTV`
    """
    def __init__(self, input_format='SMTVC', output_format='SMCTV'):
        self.input_format = input_format
        self.out_format = output_format

    def __call__(self, results, *args, **kwargs):
        keypoint = results['keypoint']
        keypoint = np.einsum(f'{self.input_format}->{self.out_format}', keypoint)
        results['keypoint'] = keypoint
        results['format'] = self.out_format
        return results


class Collect:
    def __init__(self, keys):
        self.keys = keys

    def __call__(self, results,  *args, **kwargs):
        data = {}
        for key in self.keys:
            data[key] = results[key]
        return data


class ToTensor:
    def __init__(self, keys):
        self.keys = keys

    def __call__(self, results, *args, **kwargs):
        for key in self.keys:
            results[key] = self.to_tensor(results[key])
        return results

    @staticmethod
    def to_tensor(data):
        if isinstance(data, torch.Tensor):
            return data
        if isinstance(data, np.ndarray):
            return torch.from_numpy(data).float()
        if isinstance(data, int):
            return torch.LongTensor([data])
        if isinstance(data, float):
            return torch.FloatTensor([data])
        raise TypeError(f'type {type(data)} cannot be converted to tensor.')

